﻿using System;
using System.Collections.Generic;
using System.Text;

namespace MLForgeSharp.Models.SupervisedLearningModels.DecisionTree
{
    /// <summary>
    /// 分类回归决策树
    /// </summary>
    public class CARTModel
    {
        private class TreeNode
        {
            public string FeatureName { get; set; } // 分裂特征
            public double SplitValue { get; set; } // 分裂值（用于连续特征）
            public TreeNode LeftChild { get; set; } // 左子节点
            public TreeNode RightChild { get; set; } // 右子节点
            public string Label { get; set; } // 叶节点的标签（分类任务）
            public double Value { get; set; } // 叶节点的值（回归任务）
        }

        private TreeNode root; // 决策树的根节点
        private bool isClassification; // 是否为分类任务

        public CARTModel(bool isClassification = true)
        {
            this.isClassification = isClassification;
            root = new TreeNode();
        }

        // 计算基尼不纯度
        private double CalculateGiniImpurity(int[] counts)
        {
            double gini = 1.0;
            int total = 0;
            foreach (int count in counts)
            {
                total += count;
            }
            foreach (int count in counts)
            {
                if (total > 0)
                {
                    double probability = (double)count / total;
                    gini -= probability * probability;
                }
            }
            return gini;
        }

        // 计算平方误差
        private double CalculateSquaredError(List<double> values)
        {
            double mean = 0.0;
            foreach (double value in values)
            {
                mean += value;
            }
            mean /= values.Count;

            double squaredError = 0.0;
            foreach (double value in values)
            {
                squaredError += Math.Pow(value - mean, 2);
            }
            return squaredError;
        }

        // 构建决策树
        private TreeNode BuildTree(List<double[]> data, List<string> featureNames, int labelIndex)
        {
            if (data.Count == 0)
            {
                return new TreeNode { Label = "Unknown", Value = 0.0 };
            }

            if (isClassification)
            {
                int[] labelCounts = GetLabelCounts(data, labelIndex);
                if (labelCounts.Length == 1)
                {
                    return new TreeNode { Label = data[0][labelIndex].ToString() };
                }
            }
            else
            {
                if (AllValuesEqual(data, labelIndex))
                {
                    return new TreeNode { Value = data[0][labelIndex] };
                }
            }

            if (featureNames.Count == 0)
            {
                if (isClassification)
                {
                    return new TreeNode { Label = GetMajorityLabel(data, labelIndex) };
                }
                else
                {
                    return new TreeNode { Value = GetMeanValue(data, labelIndex) };
                }
            }

            int bestFeatureIndex = 0;
            double bestSplitValue = 0.0;
            double bestImpurity = double.MaxValue;

            for (int i = 0; i < featureNames.Count; i++)
            {
                List<double> featureValues = new List<double>();
                foreach (double[] row in data)
                {
                    featureValues.Add(row[i]);
                }
                featureValues.Sort();

                for (int j = 1; j < featureValues.Count; j++)
                {
                    double splitValue = (featureValues[j - 1] + featureValues[j]) / 2;
                    List<double[]> leftSubset = new List<double[]>();
                    List<double[]> rightSubset = new List<double[]>();

                    foreach (double[] row in data)
                    {
                        if (row[i] <= splitValue)
                        {
                            leftSubset.Add(row);
                        }
                        else
                        {
                            rightSubset.Add(row);
                        }
                    }

                    double impurity = 0.0;
                    if (isClassification)
                    {
                        impurity = CalculateGiniImpurity(GetLabelCounts(leftSubset, labelIndex)) * leftSubset.Count / data.Count +
                                   CalculateGiniImpurity(GetLabelCounts(rightSubset, labelIndex)) * rightSubset.Count / data.Count;
                    }
                    else
                    {
                        impurity = CalculateSquaredError(GetLabelValues(leftSubset, labelIndex)) +
                                   CalculateSquaredError(GetLabelValues(rightSubset, labelIndex));
                    }

                    if (impurity < bestImpurity)
                    {
                        bestImpurity = impurity;
                        bestFeatureIndex = i;
                        bestSplitValue = splitValue;
                    }
                }
            }

            TreeNode node = new TreeNode
            {
                FeatureName = featureNames[bestFeatureIndex],
                SplitValue = bestSplitValue
            };

            List<double[]> leftData = new List<double[]>();
            List<double[]> rightData = new List<double[]>();
            foreach (double[] row in data)
            {
                if (row[bestFeatureIndex] <= bestSplitValue)
                {
                    leftData.Add(row);
                }
                else
                {
                    rightData.Add(row);
                }
            }

            List<string> remainingFeatures = new List<string>(featureNames);
            remainingFeatures.RemoveAt(bestFeatureIndex);

            node.LeftChild = BuildTree(leftData, remainingFeatures, labelIndex);
            node.RightChild = BuildTree(rightData, remainingFeatures, labelIndex);

            return node;
        }

        // 获取标签的计数
        private int[] GetLabelCounts(List<double[]> data, int labelIndex)
        {
            Dictionary<double, int> labelCounts = new Dictionary<double, int>();
            foreach (double[] row in data)
            {
                double label = row[labelIndex];
                if (!labelCounts.ContainsKey(label))
                {
                    labelCounts[label] = 0;
                }
                labelCounts[label]++;
            }

            int[] counts = new int[labelCounts.Count];
            int i = 0;
            foreach (int count in labelCounts.Values)
            {
                counts[i++] = count;
            }
            return counts;
        }

        // 获取标签的值（回归任务）
        private List<double> GetLabelValues(List<double[]> data, int labelIndex)
        {
            List<double> values = new List<double>();
            foreach (double[] row in data)
            {
                values.Add(row[labelIndex]);
            }
            return values;
        }

        // 检查所有值是否相等
        private bool AllValuesEqual(List<double[]> data, int labelIndex)
        {
            double firstValue = data[0][labelIndex];
            foreach (double[] row in data)
            {
                if (row[labelIndex] != firstValue)
                {
                    return false;
                }
            }
            return true;
        }

        // 获取多数标签
        private string GetMajorityLabel(List<double[]> data, int labelIndex)
        {
            Dictionary<double, int> labelCounts = new Dictionary<double, int>();
            foreach (double[] row in data)
            {
                double label = row[labelIndex];
                if (!labelCounts.ContainsKey(label))
                {
                    labelCounts[label] = 0;
                }
                labelCounts[label]++;
            }

            double majorityLabel = 0.0;
            int maxCount = 0;
            foreach (var kvp in labelCounts)
            {
                if (kvp.Value > maxCount)
                {
                    maxCount = kvp.Value;
                    majorityLabel = kvp.Key;
                }
            }
            return majorityLabel.ToString();
        }

        // 获取均值（回归任务）
        private double GetMeanValue(List<double[]> data, int labelIndex)
        {
            double sum = 0.0;
            foreach (double[] row in data)
            {
                sum += row[labelIndex];
            }
            return sum / data.Count;
        }

        // 训练模型
        public void Train(List<double[]> data, List<string> featureNames, int labelIndex)
        {
            root = BuildTree(data, featureNames, labelIndex);
        }

        // 预测
        public string Predict(double[] features, List<string> featureNames)
        {
            TreeNode node = root;
            while (node.Label == null && node.Value == 0.0)
            {
                int featureIndex = featureNames.IndexOf(node.FeatureName);
                if (features[featureIndex] <= node.SplitValue)
                {
                    node = node.LeftChild;
                }
                else
                {
                    node = node.RightChild;
                }
            }
            return isClassification ? node.Label : node.Value.ToString();
        }
    }

    // 示例程序
    public class CARTModelExample
    {
        public CARTModelExample()
        {
            // 示例数据（分类任务）
            List<double[]> data = new List<double[]>
        {
            new double[] { 1.0, 2.0, 0.0 },
            new double[] { 1.0, 3.0, 0.0 },
            new double[] { 2.0, 2.0, 1.0 },
            new double[] { 2.0, 3.0, 1.0 }
        };

            List<string> featureNames = new List<string> { "Feature1", "Feature2" };
            int labelIndex = 2; // 标签列索引

            // 创建模型
            CARTModel model = new CARTModel(isClassification: true);

            // 训练模型
            model.Train(data, featureNames, labelIndex);

            // 预测
            double[] testFeatures = { 1.0, 2.0 };
            string predictedLabel = model.Predict(testFeatures, featureNames);

            Console.WriteLine("预测标签: " + predictedLabel);
        }
    }
}
